# main.py
import os

from data_loader import load_data
from model_trainer import ModelTrainer
from model_predictor import ModelPredictor
from config import DATA_PATH

def train_pipeline():
    print("正在加载训练数据...")
    texts, labels = load_data(DATA_PATH["train"])

    trainer = ModelTrainer()
    trainer.train(texts, labels)

def predict_pipeline():
    print("正在加载测试数据...")
    texts, _ = load_data(DATA_PATH["test"])

    predictor = ModelPredictor()
    predictor.predict(texts)

if __name__ == "__main__":
    # 训练模型
    if os.path.exists(DATA_PATH["train"]):
        train_pipeline()
    else:
        print(f"训练文件不存在: {DATA_PATH['train']}")

    # 预测模型
    if os.path.exists(DATA_PATH["test"]):
        predict_pipeline()
    else:
        print(f"测试文件不存在: {DATA_PATH['test']}")